{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lecture 2\n",
    "基本上算是老师代码的复刻版本。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Policy Iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np \n",
    "import gym\n",
    "from gym import wrappers\n",
    "from gym.envs.registration import register"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_episode(env, policy, gamma=1.0, render=False):\n",
    "    '''\n",
    "    Run an episode and return reward.\n",
    "    '''\n",
    "    obs = env.reset()\n",
    "    total_reward = 0\n",
    "    step_idx = 0\n",
    "    while True:\n",
    "        if render:\n",
    "            # 暂时不清楚 render 是什么作用\n",
    "            env.render()\n",
    "        # env.step 返回四个值，\n",
    "        # 1.observation: 当前观测：图像的像素值 \n",
    "        # 2.reward: 当前奖励\n",
    "        # 3.gameover: True 为游戏结束\n",
    "        # 4.额外信息\n",
    "        obs, reward, done, _ = env.step(int(policy[obs]))\n",
    "        total_reward += (gamma ** step_idx * reward)\n",
    "        step_idx += 1\n",
    "        if done:\n",
    "            break\n",
    "    return total_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_policy(env, policy, gamma=1.0, n=100):\n",
    "    '''\n",
    "    evaluate a policy is to calculate the value function.\n",
    "    Using Monte-Carlo method.\n",
    "    '''\n",
    "    scores =[run_episode(env, policy, gamma, False) for _ in range(n)]\n",
    "    return np.mean(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_policy(v, gamma=1.0):\n",
    "    '''\n",
    "    Extract the policy given a value-function.\n",
    "    '''\n",
    "    policy = np.zeros(env.env.nS)\n",
    "    for s in range(env.env.nS):\n",
    "        q_sa = np.zeros(env.env.nA)\n",
    "        for a in range(env.env.nA):\n",
    "            q_sa[a] = sum([p * (r + gamma * v[s_]) for p,s_,r,_ in env.env.P[s][a]])\n",
    "        policy[s] = np.argmax(q_sa)\n",
    "    return policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_policy_v(env, policy, gamma=1.0):\n",
    "    '''\n",
    "    Iteratively evaluate the value-function under policy.\n",
    "    '''\n",
    "    v = np.zeros(env.env.nS)\n",
    "    eps = 1e-10\n",
    "    while True:\n",
    "        prev_v = np.copy(v)\n",
    "        for s in range(env.env.nS):\n",
    "            policy_a = policy[s]\n",
    "            v[s] = sum([p * (r + gamma * prev_v[s_]) for p,s_,r,_ in env.env.P[s][policy_a]])\n",
    "        if (np.sum(np.fabs(prev_v - v))) <= eps:\n",
    "            break\n",
    "    return v"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def policy_iteration(env, gamma=1.0):\n",
    "    '''\n",
    "    policy-iteration algorithm.\n",
    "    '''\n",
    "    # 随机初始化 policy\n",
    "    policy = np.random.choice(env.env.nA, size=(env.env.nS))\n",
    "    max_iteration = 200000\n",
    "    gamma = 1.0\n",
    "    for i in range(max_iteration):\n",
    "        old_policy_v = compute_policy_v(env,policy,gamma)\n",
    "        new_policy = extract_policy(old_policy_v,gamma)\n",
    "        if (np.all(policy == new_policy)):\n",
    "            print('Policy-Iteration converged at step %d.' % (i+1))\n",
    "            break\n",
    "        policy = new_policy\n",
    "    return policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Policy-Iteration converged at step 4.\n",
      "Average scores =  0.74\n"
     ]
    }
   ],
   "source": [
    "env_name  = 'FrozenLake-v0' # 'FrozenLake8x8-v0'\n",
    "env = gym.make(env_name)\n",
    "\n",
    "optimal_policy = policy_iteration(env, gamma = 1.0)\n",
    "scores = evaluate_policy(env, optimal_policy, gamma = 1.0)\n",
    "print('Average scores = ', np.mean(scores))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "torch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
